#include <iostream>


// the size of input must be 2**N
__global__ void global_reduce_kernel(float* d_out, float* d_in)
{
    int gid = threadIdx.x + blockDim.x * blockIdx.x;
    int tid = threadIdx.x;

    for(unsigned int s = blockDim.x / 2; s > 0; s >>= 1)
    {
        if(tid < s){
            d_in[gid] += d_in[gid + s];
        }
        __syncthreads();
    } 

    // 仅合并0线程的结果
    if(tid ==0){
        d_out[blockIdx.x] = d_in[gid];
        printf("%d, %f\n", blockIdx.x, d_out[blockIdx.x]);
    }

}

__global__ void shared_reduce_kernel(float* d_out, float* d_in, int size)
{
    extern __shared__ float sdata[];

    int tid = threadIdx.x;
    int gid = tid + blockDim.x * blockIdx.x;
    
    // copy data to shared data
    if(gid >= size) sdata[tid] = 0;
    else sdata[tid] = d_in[gid];
    __syncthreads();

    // reduce
    for(unsigned int s = blockDim.x / 2; s > 0; s >>= 1)
    {
        if(tid < s){
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    // copy shared data to output data
    if(tid==0){
        d_out[blockIdx.x] = sdata[0];
    }

}


int main(int argc, char** argv)
{

    // get cpu data
    const int size = 800;
    float h_x[size];
    float sum(0);
    for(int i = 0; i < size; ++i){
        h_x[i] = i;
        sum += i;
    }

    // get gpu data
    float* d_x;
    cudaMalloc((void**)&d_x, sizeof(float)*size);
    cudaMemcpy(d_x, h_x, sizeof(float)*size, cudaMemcpyHostToDevice);

    int block_size = 13;
    float* d_y;
    cudaMalloc((void**)&d_y, sizeof(float) * block_size);

    // run kernel
    dim3 Db(64);
    // dim3 Dg(size / 64 + 1);
    shared_reduce_kernel<<<block_size, Db, sizeof(float)* Db.x>>>(d_y, d_x, size);  

    // gpu to cpu
    float* y = new float[block_size];
    cudaMemcpy(y, d_y, sizeof(float)*block_size, cudaMemcpyDeviceToHost);

    int gpu_sum(0);
    for(int i = 0; i < block_size; ++i){
        gpu_sum += y[i];
    }

    std::cout << "the gpu result is:" << gpu_sum << std::endl;
    std::cout << "the correct result is:" << sum << std::endl;

    cudaFree(d_x);
    cudaFree(d_y);
    delete[] y;
    return 0;
}